from celery_config import celery_app
from ultralytics import YOLO
import cv2
import os
from datetime import datetime
from database import SessionLocal
from models import Task, TASK_STATUS_PROCESSING, TASK_STATUS_COMPLETED, TASK_STATUS_FAILED, DetectionDetail, VehicleDetail, TaskStatistics, History, DetectionRecord, UserStatistics
from sqlalchemy.orm import Session
import numpy as np
import logging
import json
import subprocess
import sys
import ssl
from typing import Dict
from fastapi import HTTPException

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 获取当前文件所在目录的绝对路径
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# 构建模型目录路径
MODEL_DIR = os.path.join(os.path.dirname(BASE_DIR), "Models-----YOLO11")
# 构建上传文件目录路径
UPLOAD_DIR = os.path.join(os.path.dirname(BASE_DIR), "uploads")

# 加载模型
MODEL_PATH = os.path.join(MODEL_DIR, "yolo11n.pt")
yolo_model = YOLO(MODEL_PATH)

def process_image(task_id, input_path, conf, iou, max_det, db, user_id):
    """处理图片文件"""
    logger.info(f"Processing image: {input_path}")
    try:
        # 确保输入文件存在
        if not os.path.exists(input_path):
            raise Exception(f"输入文件不存在: {input_path}")

        # 生成输出文件路径
        output_path = input_path.replace("input_", "output_")
        logger.info(f"Output path: {output_path}")

        # 运行检测
        results = yolo_model(
            input_path,
            conf=conf,
            iou=iou,
            max_det=max_det
        )
        result = results[0]
        boxes = result.boxes

        # 只保留人和车的类别
        target_labels = ['person', 'car', 'bus', 'truck', 'motorcycle']
        filtered_boxes = []
        for box in boxes:
            cls = int(box.cls[0])
            label = yolo_model.model.names[cls]
            if label in target_labels:
                filtered_boxes.append(box)

        # 读取原始图片
        img = cv2.imread(input_path)
        if img is None:
            raise Exception(f"无法读取图片: {input_path}")

        # 画框
        for box in filtered_boxes:
            xyxy = box.xyxy[0].cpu().numpy().astype(int)
            cls = int(box.cls[0])
            conf_val = float(box.conf[0])
            label = yolo_model.model.names[cls]
            color = (0, 255, 0) if label == 'person' else (255, 0, 0)
            cv2.rectangle(img, (xyxy[0], xyxy[1]), (xyxy[2], xyxy[3]), color, 2)
            cv2.putText(img, f"{label} {conf_val:.2f}", (xyxy[0], xyxy[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)

        # 保存处理后的图片
        success = cv2.imwrite(output_path, img)
        if not success:
            raise Exception(f"无法保存处理后的图片: {output_path}")

        # 验证输出文件是否成功保存
        if not os.path.exists(output_path):
            raise Exception(f"输出文件未成功保存: {output_path}")

        logger.info(f"Successfully saved processed image to: {output_path}")

        # 统计数量
        pedestrian_count = sum(1 for box in filtered_boxes if yolo_model.model.names[int(box.cls[0])] == 'person')
        vehicle_count = sum(1 for box in filtered_boxes if yolo_model.model.names[int(box.cls[0])] in ['car', 'bus', 'truck', 'motorcycle'])
        vehicle_types = {}
        for box in filtered_boxes:
            label = yolo_model.model.names[int(box.cls[0])]
            if label in ['car', 'bus', 'truck', 'motorcycle']:
                vehicle_types[label] = vehicle_types.get(label, 0) + 1

        logger.info(f"Detected: Pedestrians={pedestrian_count}, Vehicles={vehicle_count}, Vehicle Types={vehicle_types}")

        # 写入 DetectionDetail
        for box in filtered_boxes:
            cls = int(box.cls[0])
            label = yolo_model.model.names[cls]
            conf_val = float(box.conf[0])
            detection_detail = DetectionDetail(
                task_id=task_id,
                object_type=label,
                confidence=conf_val,
                frame_start=0,
                frame_end=0
            )
            db.add(detection_detail)
        logger.info(f"Added {len(filtered_boxes)} DetectionDetail records to session")
        db.commit()
        logger.info("DetectionDetail records committed")

        # 写入 VehicleDetail
        for vehicle_type, count in vehicle_types.items():
            vehicle_detail = VehicleDetail(
                task_id=task_id,
                vehicle_type=vehicle_type,
                count=count
            )
            db.add(vehicle_detail)
        logger.info(f"Added {len(vehicle_types)} VehicleDetail records to session")
        db.commit()
        logger.info("VehicleDetail records committed")

        # 写入 DetectionRecord
        detection_results = []
        for box in filtered_boxes:
            cls = int(box.cls[0])
            label = yolo_model.model.names[cls]
            conf_val = float(box.conf[0])
            detection_results.append({
                "object_type": label,
                "confidence": conf_val
            })
        detection_record = DetectionRecord(
            user_id=user_id,
            image_path=input_path,
            confidence_threshold=conf,
            iou_threshold=iou,
            max_detections=max_det,
            detection_results=json.dumps(detection_results)
        )
        db.add(detection_record)
        logger.info("Added DetectionRecord to session")
        db.commit()
        logger.info("DetectionRecord committed")

        # 返回结果
        result = {
            'success': True,
            'output_path': output_path,
            'pedestrian_count': pedestrian_count,
            'vehicle_count': vehicle_count,
            'vehicle_types': vehicle_types,
            'detection_details': [{
                "object_type": yolo_model.model.names[int(box.cls[0])],
                "confidence": float(box.conf[0]),
                "frame_start": 0,
                "frame_end": 0
            } for box in filtered_boxes],
            'detection_record': {
                "user_id": user_id,
                "image_path": input_path,
                "confidence_threshold": conf,
                "iou_threshold": iou,
                "max_detections": max_det,
                "detection_results": detection_results
            }
        }
        logger.info(f"Processing result: {result}")
        return result
    except Exception as e:
        logger.error(f"Error processing image: {str(e)}")
        return {
            'success': False,
            'error': str(e)
        }

def process_video(task_id, input_path, conf, iou, max_det, db, user_id):
    """处理视频文件"""
    logger.info(f"Processing video: {input_path}")
    try:
        # 直接使用基本检测模式
        return process_video_basic(task_id, input_path, conf, iou, max_det, db, user_id)
    except Exception as e:
        logger.error(f"Error processing video: {str(e)}")
        return {
            'success': False,
            'error': str(e)
        }

def process_video_basic(task_id, input_path, conf, iou, max_det, db, user_id):
    """基本的视频处理函数，不使用 ObjectCounter"""
    logger.info("Using basic video processing mode")
    try:
        cap = cv2.VideoCapture(input_path)
        if not cap.isOpened():
            raise Exception("无法打开视频文件")
        
        # 获取视频信息
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        
        # 创建输出文件
        output_path = input_path.replace("input_", "output_")
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        # 初始化计数器
        max_pedestrian_count = 0  # 最大人数
        max_car_count = 0  # 最大车辆数
        frame_count = 0
        detection_results = []
        frame_detection_details = []
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            
            # 处理每一帧
            results = yolo_model(frame, conf=conf, iou=iou, max_det=max_det)
            result = results[0]
            
            # 当前帧的计数
            current_pedestrian_count = 0
            current_car_count = 0
            
            # 获取检测结果
            for box in result.boxes:
                cls = int(box.cls[0])
                label = yolo_model.model.names[cls]
                
                # 只处理人和车
                if label not in ['person', 'car']:
                    continue
                
                # 获取边界框坐标
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                conf_val = float(box.conf[0])
                
                # 更新当前帧的计数
                if label == 'person':
                    current_pedestrian_count += 1
                else:  # car
                    current_car_count += 1
                
                # 记录检测结果
                detection_results.append({
                    "object_type": label,
                    "confidence": conf_val,
                    "frame": frame_count
                })
                
                # Store detection details to be returned
                frame_detection_details.append({
                    "object_type": label,
                    "confidence": conf_val,
                    "frame_start": frame_count,
                    "frame_end": frame_count # For videos, frame_start and frame_end are the same for basic mode
                })
            
            # 更新最大计数
            max_pedestrian_count = max(max_pedestrian_count, current_pedestrian_count)
            max_car_count = max(max_car_count, current_car_count)
            
            db.commit()
            
            # 在帧上绘制检测结果
            for box in result.boxes:
                cls = int(box.cls[0])
                label = yolo_model.model.names[cls]
                
                # 只处理人和车
                if label not in ['person', 'car']:
                    continue
                
                # 获取边界框坐标
                x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
                conf_val = float(box.conf[0])
                
                # 设置颜色
                color = (0, 255, 0) if label == 'person' else (255, 0, 0)
                
                # 绘制边界框
                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                
                # 添加标签
                label_text = f"{label} {conf_val:.2f}"
                cv2.putText(frame, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
            
            # 写入输出视频
            out.write(frame)
            frame_count += 1
        
        # 释放资源
        cap.release()
        out.release()
        
        # Prepare detection results for DetectionRecord
        detection_record_results = []
        # For basic video processing, DetectionRecord might store a summary or per-frame data.
        # For simplicity, let's just store the detection results list gathered during frame processing.
        # A more detailed implementation might aggregate this.
        detection_record_results = detection_results # Reuse the list collected during frame processing
        
        # Return results
        result = {
            'success': True,
            'output_path': output_path,
            'pedestrian_count': max_pedestrian_count, # Return max counts for video
            'vehicle_count': max_car_count,      # Return max counts for video
            'vehicle_types': {}, # Basic video processing doesn't aggregate vehicle types easily
            'detection_details': frame_detection_details, # Return collected details
            'detection_record': {
                "user_id": user_id,
                "image_path": input_path,
                "confidence_threshold": conf,
                "iou_threshold": iou,
                "max_detections": max_det,
                "detection_results": detection_record_results # Use the collected results
            }
        }
        logger.info(f"Processing result: {result}")
        return result
    except Exception as e:
        logger.error(f"Error processing video: {str(e)}")
        return {
            'success': False,
            'error': str(e)
        }

def calculate_iou(box1, box2):
    """计算两个边界框的IOU"""
    # 获取边界框坐标
    x1_1, y1_1, x2_1, y2_1 = box1
    x1_2, y1_2, x2_2, y2_2 = box2
    
    # 计算交集区域
    x1_i = max(x1_1, x1_2)
    y1_i = max(y1_1, y1_2)
    x2_i = min(x2_1, x2_2)
    y2_i = min(y2_1, y2_2)
    
    # 计算交集面积
    if x2_i < x1_i or y2_i < y1_i:
        return 0.0
    
    intersection = (x2_i - x1_i) * (y2_i - y1_i)
    
    # 计算两个边界框的面积
    box1_area = (x2_1 - x1_1) * (y2_1 - y1_1)
    box2_area = (x2_2 - x1_2) * (y2_2 - y1_2)
    
    # 计算并集面积
    union = box1_area + box2_area - intersection
    
    # 计算IOU
    iou = intersection / union if union > 0 else 0
    
    return iou

@celery_app.task
def process_task(task_id: str, input_path: str, conf: float, iou: float, max_det: int):
    """处理检测任务"""
    logger.info(f"\n{'='*50}")
    logger.info(f"[{datetime.now()}] Starting task processing")
    logger.info(f"Task ID: {task_id}")
    logger.info(f"Input file: {input_path}")
    logger.info(f"Parameters: conf={conf}, iou={iou}, max_det={max_det}")
    
    # 创建数据库会话
    db = SessionLocal()
    try:
        # 获取任务信息
        task = db.query(Task).filter(Task.id == task_id).first()
        if not task:
            logger.error(f"Error: Task {task_id} not found")
            return
        
        # 验证输入文件
        if not os.path.exists(input_path):
            error_msg = f"输入文件不存在: {input_path}"
            logger.error(error_msg)
            task.status = TASK_STATUS_FAILED
            task.error = error_msg
            db.commit()
            return
        
        # 更新任务状态为处理中
        task.status = TASK_STATUS_PROCESSING
        db.commit()
        logger.info(f"Task status updated to PROCESSING")
        
        # 确定媒体类型
        media_type = "image" if input_path.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif')) else "video"
        task.media_type = media_type
        db.commit()
        logger.info(f"Media type set to: {media_type}")
        
        # 记录开始时间
        start_time = datetime.now()
        
        # 处理文件
        if media_type == "image":
            # 处理图片
            result = process_image(task_id, input_path, conf, iou, max_det, db, task.user_id)
        else:
            # 处理视频
            result = process_video(task_id, input_path, conf, iou, max_det, db, task.user_id)
        
        # 计算处理时间
        processing_time = (datetime.now() - start_time).total_seconds()
        logger.info(f"Processing completed in {processing_time:.2f} seconds")
        
        # 更新任务状态和结果
        if result.get('success', False):
            output_path = result.get('output_path')
            pedestrian_count = result.get('pedestrian_count', 0)
            vehicle_count = result.get('vehicle_count', 0)
            vehicle_types = result.get('vehicle_types', {})
            detection_details_data = result.get('detection_details', [])
            detection_record_data = result.get('detection_record')

            logger.info(f"Processing result output_path: {output_path}")
            logger.info(f"Detected counts: Pedestrians={pedestrian_count}, Vehicles={vehicle_count}")
            logger.info(f"Vehicle types: {vehicle_types}")
            logger.info(f"Number of detection details: {len(detection_details_data)}")
            logger.info(f"Detection record data present: {detection_record_data is not None}")

            # 验证输出文件
            if not output_path or not os.path.exists(output_path):
                error_msg = f"处理结果文件不存在或路径未设置: {output_path}"
                logger.error(error_msg)
                task.status = TASK_STATUS_FAILED
                task.error = error_msg
                db.commit() # Commit failure status
                return

            # Update task information
            task.status = TASK_STATUS_COMPLETED
            task.output_file = output_path
            task.pedestrian_count = pedestrian_count
            task.vehicle_count = vehicle_count
            task.processing_time = processing_time
            # db.commit() # Moved commit to the end
            logger.info(f"Task status updated to COMPLETED, output_file set to {output_path}")

            # Create History record
            history = History(
                user_id=task.user_id,
                input_file=input_path,
                output_file=output_path,
                pedestrian_count=pedestrian_count,
                vehicle_count=vehicle_count,
                confidence_threshold=conf,
                iou_threshold=iou,
                max_detections=max_det,
                media_type=media_type
            )
            db.add(history)
            db.flush()  # Get history.id before committing
            logger.info(f"History record created for task {task_id} with ID {history.id}")

            # Add VehicleDetail records
            for vehicle_type, count in vehicle_types.items():
                vehicle_detail = VehicleDetail(
                    task_id=task_id,
                    vehicle_type=vehicle_type,
                    count=count
                )
                db.add(vehicle_detail)
            logger.info(f"Added {len(vehicle_types)} VehicleDetail records")

            # Add DetectionDetail records
            for detail_data in detection_details_data:
                detection_detail = DetectionDetail(
                    task_id=task_id,
                    object_type=detail_data.get('object_type'),
                    confidence=detail_data.get('confidence'),
                    frame_start=detail_data.get('frame_start'),
                    frame_end=detail_data.get('frame_end')
                )
                db.add(detection_detail)
            logger.info(f"Added {len(detection_details_data)} DetectionDetail records")

            # Add DetectionRecord
            if detection_record_data:
                detection_record = DetectionRecord(
                    user_id=detection_record_data.get('user_id'),
                    image_path=detection_record_data.get('image_path'),
                    confidence_threshold=detection_record_data.get('confidence_threshold'),
                    iou_threshold=detection_record_data.get('iou_threshold'),
                    max_detections=detection_record_data.get('max_detections'),
                    detection_results=json.dumps(detection_record_data.get('detection_results', []))
                )
                db.add(detection_record)
                logger.info("Added DetectionRecord")

            # Update user statistics
            update_user_statistics(task.user_id, media_type, result, db) # This function already commits

            # Commit all changes for this task
            db.commit()
            logger.info("All task-related database changes committed successfully.")

        else:
            error_msg = result.get('error', '处理失败')
            logger.error(f"Task failed: {error_msg}")
            task.status = TASK_STATUS_FAILED
            task.error = error_msg
            db.commit()
        
        logger.info(f"{'='*50}\n")
        return result
        
    except Exception as e:
        error_msg = f"处理任务时发生错误: {str(e)}"
        logger.error(error_msg)
        if task:
            task.status = TASK_STATUS_FAILED
            task.error = error_msg
            db.commit()
        raise
    finally:
        db.close()

def update_user_statistics(user_id: int, media_type: str, result: dict, db: Session):
    """更新用户统计信息"""
    try:
        # 获取用户统计信息
        stats = db.query(UserStatistics).filter(UserStatistics.user_id == user_id).first()
        
        # 如果不存在，创建新的统计记录
        if not stats:
            stats = UserStatistics(
                user_id=user_id,
                total_tasks=0,
                completed_tasks=0,
                failed_tasks=0,
                total_images=0,
                total_videos=0,
                total_pedestrians=0,
                total_vehicles=0
            )
            db.add(stats)
        
        # 更新统计信息
        stats.total_tasks = (stats.total_tasks or 0) + 1
        stats.completed_tasks = (stats.completed_tasks or 0) + 1
        
        # 根据媒体类型更新计数
        if media_type == 'image':
            stats.total_images = (stats.total_images or 0) + 1
        elif media_type == 'video':
            stats.total_videos = (stats.total_videos or 0) + 1
        
        # 更新检测结果统计
        stats.total_pedestrians = (stats.total_pedestrians or 0) + (result.get('pedestrian_count', 0) or 0)
        stats.total_vehicles = (stats.total_vehicles or 0) + (result.get('vehicle_count', 0) or 0)
        
        # 更新最后更新时间
        stats.last_updated = datetime.now()
        
        # 提交更改
        db.commit()
        logger.info(f"User statistics updated successfully for user {user_id}")
        
    except Exception as e:
        logger.error(f"Error updating user statistics: {str(e)}")
        db.rollback()
        raise

def get_task_result(task_id: str) -> Dict:
    """获取任务结果"""
    task = Task.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    
    result = {
        "task_id": task.task_id,
        "status": task.status,
        "result": task.result,
        "error": task.error,
        "confidence_threshold": task.confidence_threshold,
        "iou_threshold": task.iou_threshold,
        "max_detections": task.max_detections
    }
    return result 